import sys
import torch
import torch.nn.functional as F
from pathlib import Path

# 添加本地路径到 sys.path 的最前面，确保优先使用本地模块
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

# 导入本地自定义模块并注册到 ultralytics
import nn.modules.conv as local_conv
import nn.modules.block as local_block
import nn.tasks as local_tasks

# 将自定义模块和parse_model注入到 ultralytics 中
try:
    import ultralytics.nn.tasks as tasks_module
    # 注册自定义模块到 globals
    tasks_module.GSConv = local_conv.GSConv
    tasks_module.GAM = local_conv.GAM
    tasks_module.AdaptiveGAM = local_conv.AdaptiveGAM  # 新增：动态自适应注意力
    tasks_module.VoVGSCSP = local_block.VoVGSCSP
    tasks_module.SAVPE = local_block.SAVPE
    tasks_module.ProgressiveFusionModule = local_block.ProgressiveFusionModule  # 新增：渐进式融合
    # 重要：替换 parse_model 函数以支持自定义模块
    tasks_module.parse_model = local_tasks.parse_model

    def patched_predict_once(self, x, profile=False, visualize=False, embed=None):
        y, dt, embeddings = [], [], []
        for m in self.model:
            if m.f != -1:  # if not from previous layer
                if isinstance(m.f, int):
                    x = y[m.f]
                else:
                    inputs = [x if j == -1 else y[j] for j in m.f]
                    x = inputs[0] if len(inputs) == 1 else inputs
            if profile:
                self._profile_one_layer(m, x, dt)
            x = m(x)
            y.append(x if m.i in self.save else None)
            if visualize:
                tasks_module.feature_visualization(x, m.type, m.i, save_dir=visualize)
            if embed and m.i in embed:
                embeddings.append(F.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1))
                if m.i == max(embed):
                    return torch.unbind(torch.cat(embeddings, 1), dim=0)
        return x

    tasks_module.BaseModel._predict_once = patched_predict_once

    print("✅ 自定义模块已成功注册:")
    print("   - GSConv, GAM, AdaptiveGAM (动态自适应注意力)")
    print("   - VoVGSCSP, SAVPE")
    print("   - ProgressiveFusionModule (渐进式特征融合)")
    print("✅ 自定义 parse_model 已注册")
except Exception as e:
    print(f"⚠️ 注册自定义模块时出错: {e}")

# 注册改进的损失函数
try:
    import utils.loss as local_loss
    import ultralytics.utils.loss as ul_loss
    # 注册WBboxLoss
    ul_loss.WBboxLoss = local_loss.WBboxLoss
    print("✅ 改进损失函数已注册: WBboxLoss (WIoU)")
except Exception as e:
    print(f"⚠️ 注册损失函数时出错: {e}")

from ultralytics import YOLO

if __name__ == '__main__':
    # 使用自定义模型配置文件
    model = YOLO("yolov8-custom-advanced.yaml")
    
    print("\n" + "="*60)
    print("🚀 YOLOv8 Custom Training with Advanced Features")
    print("="*60)
    print("创新点1: WIoU Loss (动态质量权重)")
    print("创新点2: 类别平衡Focal Loss (自适应γ)")
    print("创新点3: AdaptiveGAM (动态注意力机制)")
    print("创新点4: Progressive Feature Fusion (渐进式融合)")
    print("="*60 + "\n")
    
    # 数据集配置文件
    model.train(data=r"mydata.yaml",
                epochs=100,
                batch=1,
                workers=0,
                patience=80,
                imgsz=640,
                device='cpu',  # 使用 CPU，如果有 GPU 可以改为 0
                )
